package com.booter.gateway.filter;

import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.filter.factory.rewrite.ModifyRequestBodyGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.factory.rewrite.ModifyRequestBodyGatewayFilterFactory.Config;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.List;

import static org.springframework.cloud.gateway.support.GatewayToStringStyler.filterToStringCreator;

/***
 * 用于修改请求的内容
 * 设计请求的contentType
 */
@Component
public class CustomModifyRequestBodyGatewayFilterFactory extends AbstractGatewayFilterFactory<Config> {


    private final List<HttpMessageReader<?>> messageReaders;

    public CustomModifyRequestBodyGatewayFilterFactory() {
        super(Config.class);
        this.messageReaders = HandlerStrategies.withDefaults().messageReaders();
    }

    public CustomModifyRequestBodyGatewayFilterFactory(
            List<HttpMessageReader<?>> messageReaders) {
        super(Config.class);
        this.messageReaders = messageReaders;
    }

    private Config getConfig() {
        Config conf = new Config();
        conf.setContentType(MediaType.APPLICATION_JSON_VALUE);
        conf.setInClass(String.class);
        conf.setOutClass(String.class);
        //执行body修改动作
        conf.setRewriteFunction(new CustomRewriteFunction());
        return conf;
    }


    @Override
    @SuppressWarnings("unchecked")
    public GatewayFilter apply(ModifyRequestBodyGatewayFilterFactory.Config config) {
        Config conf = this.getConfig();
        return new GatewayFilter() {
            @Override
            public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
                Class inClass = conf.getInClass();
                ServerRequest serverRequest = ServerRequest.create(exchange, CustomModifyRequestBodyGatewayFilterFactory.this.messageReaders);
                // TODO: flux or mono
                Mono<?> modifiedBody = serverRequest.bodyToMono(inClass)
                        .flatMap(originalBody -> conf.getRewriteFunction().apply(exchange, originalBody))
                        .switchIfEmpty(Mono.defer(() -> (Mono) conf.getRewriteFunction().apply(exchange, null)));

                BodyInserter bodyInserter = BodyInserters.fromPublisher(modifiedBody, conf.getOutClass());
                HttpHeaders headers = new HttpHeaders();
                headers.putAll(exchange.getRequest().getHeaders());

                headers.remove(HttpHeaders.CONTENT_LENGTH);

                // if the body is changing content types, set it here, to the bodyInserter
                // will know about it
                if (conf.getContentType() != null) {
                    headers.set(HttpHeaders.CONTENT_TYPE, conf.getContentType());
                }
                CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);
                return bodyInserter.insert(outputMessage, new BodyInserterContext())
                        // .log("modify_request", Level.INFO)
                        .then(Mono.defer(() -> {
                            ServerHttpRequest decorator = CustomModifyRequestBodyGatewayFilterFactory.this.decorate(exchange, headers, outputMessage);
                            return chain.filter(exchange.mutate().request(decorator).build());
                        }));
            }

            @Override
            public String toString() {
                return filterToStringCreator(this)
                        .append("Content type", conf.getContentType())
                        .append("In class", conf.getInClass())
                        .append("Out class", conf.getOutClass()).toString();
            }
        };
    }

    ServerHttpRequestDecorator decorate(ServerWebExchange exchange, HttpHeaders headers,
                                        CachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public HttpHeaders getHeaders() {
                long contentLength = headers.getContentLength();
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll(headers);
                if (contentLength > 0) {
                    httpHeaders.setContentLength(contentLength);
                } else {
                    // TODO: this causes a 'HTTP/1.1 411 Length Required' // on
                    // httpbin.org
                    httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                }
                return httpHeaders;
            }

            @Override
            public Flux<DataBuffer> getBody() {
                return outputMessage.getBody();
            }
        };
    }


}
